# import cv2
# import os
# from collections import defaultdict
#
# # 数据集路径和标注文件路径
# image_folder = r'E:\b-s\yolov10\yolov10-main\dataset\valid\images'
# label_folder = r'E:\b-s\yolov10\yolov10-main\dataset\valid\labels'
#
# # 创建一个字典来统计每个标签的数量
# label_counts = defaultdict(int)
#
# for image_file in os.listdir(image_folder):
#     # 读取图像
#     img = cv2.imread(os.path.join(image_folder, image_file))
#
#     if img is None:
#         print(f"Error loading image: {image_file}")
#         continue
#
#     height, width, _ = img.shape
#
#     # 读取对应的标注文件
#     label_file = os.path.splitext(image_file)[0] + '.txt'
#     with open(os.path.join(label_folder, label_file), 'r') as f:
#         for line in f.readlines():
#             parts = line.strip().split()
#             if len(parts) != 5:
#                 print(f"Invalid format in label file: {label_file}")
#                 continue
#
#             class_id, x_center, y_center, w, h = map(float, parts)
#             # 更新标签计数
#             label_counts[int(class_id)] += 1
#
#             # 转换为像素坐标
#             x1 = int((x_center - w / 2) * width)
#             y1 = int((y_center - h / 2) * height)
#             x2 = int((x_center + w / 2) * width)
#             y2 = int((y_center + h / 2) * height)
#
#             # 绘制矩形框
#             color = (255, 0, 0)  # Red color
#             cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
#             cv2.putText(img, str(int(class_id)), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
#
#     # 显示图像
#     cv2.imshow('Image', img)
#     cv2.waitKey(1)  # 每张图像显示2秒（2000毫秒）
#
# # 输出统计结果
# print("标签统计结果:")
# for class_id, count in label_counts.items():
#     print(f"类 {class_id}: {count} 次")
#
# cv2.destroyAllWindows()


import cv2
import torch
from ultralytics import YOLO
import time
import os
import numpy as np


def detect_video():
    # 内置路径配置
    model_path = r"E:\b-s\yolov10\yolov10-main\runs\detect\train_v106\weights\best.pt"  # 模型路径，请替换为您的实际模型路径
    video_path = 0  # 视频路径，请替换为您的实际视频路径（或使用0表示默认摄像头）
    output_path = "output.mp4"  # 输出视频路径

    # 检测参数
    conf_threshold = 0.25  # 置信度阈值
    show_result = True  # 是否显示检测结果
    save_video = True  # 是否保存输出视频

    print("开始视频检测...")
    print(f"模型路径: {model_path}")
    print(f"视频源: {video_path}")

    # 加载YOLO模型
    try:
        model = YOLO(model_path)
        print(f"YOLOv10模型加载成功")
    except Exception as e:
        print(f"模型加载失败: {e}")
        return

    # 设置设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")

    # 打开视频源
    if video_path == "0" or video_path == 0:
        cap = cv2.VideoCapture(0)
        print("已打开默认摄像头")
    else:
        cap = cv2.VideoCapture(video_path)
        print(f"已打开视频文件: {video_path}")

    if not cap.isOpened():
        print(f"无法打开视频源")
        return

    # 获取视频信息
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print(f"视频尺寸: {width}x{height}, FPS: {fps}")

    # 设置输出视频
    out = None
    if save_video:
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        print(f"输出视频将保存至: {output_path}")

    # 开始处理视频
    frame_count = 0
    detection_times = []
    start_time = time.time()

    while True:
        ret, frame = cap.read()
        if not ret:
            print("视频读取结束或出错")
            break

        frame_count += 1

        # 记录检测开始时间
        detection_start = time.time()

        # 使用YOLO模型进行检测
        results = model.predict(
            source=frame,
            conf=conf_threshold,
            device=device,
            verbose=False
        )

        # 记录检测时间
        detection_time = time.time() - detection_start
        detection_times.append(detection_time)

        # 在帧上绘制检测结果
        annotated_frame = results[0].plot()

        # 添加性能信息
        avg_time = np.mean(detection_times[-30:]) if detection_times else 0
        fps_text = f"FPS: {1 / avg_time:.1f}" if avg_time > 0 else "FPS: N/A"
        cv2.putText(annotated_frame, fps_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        # 显示检测结果
        if show_result:
            cv2.imshow('YOLOv10 Detection', annotated_frame)

        # 保存到输出视频
        if out is not None:
            out.write(annotated_frame)

        # 按ESC键退出
        if show_result and (cv2.waitKey(1) & 0xFF == 27):
            print("用户按ESC键退出")
            break

        # 每100帧打印一次状态
        if frame_count % 100 == 0:
            elapsed = time.time() - start_time
            print(f"已处理 {frame_count} 帧, 平均速度: {frame_count / elapsed:.2f} FPS")

    # 计算处理帧率
    elapsed_time = time.time() - start_time
    if elapsed_time > 0 and frame_count > 0:
        processed_fps = frame_count / elapsed_time
        print(f"总共处理 {frame_count} 帧")
        print(f"平均处理帧率: {processed_fps:.2f} FPS")
        print(f"平均检测时间: {np.mean(detection_times) * 1000:.2f} ms/帧")

    # 释放资源
    cap.release()
    if out is not None:
        out.release()
    cv2.destroyAllWindows()
    print("视频检测完成")


if __name__ == "__main__":
    detect_video()
